import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from python.lib import json


def generate_mp_traffic_matrix(nodes=16, shard_nodes=[0, 3, 8, 13], traffic_size=0.03125):
    """
    生成MP流量矩阵
    :param nodes: 总节点数
    :param shard_nodes: 分片节点索引列表
    :param traffic_size: 单次传输数据量(MB)
    :return: numpy矩阵
    """
    matrix = np.zeros((nodes, nodes))

    # 为每个分片节点填充流量
    for i in shard_nodes:
        for j in range(nodes):
            if i != j:  # 分片节点向其他所有节点发送数据
                matrix[i][j] = traffic_size
    return matrix

# 定义allreduce流量矩阵
def generate_ring_allreduce(nodes=16, traffic_size=44):
    matrix = np.zeros((nodes, nodes))
    for i in range(nodes-1):
        matrix[i][i+1] = traffic_size
    matrix[nodes-1][0]=traffic_size
    return matrix

# 生成突发矩阵
def generate_burst_traffic(hotspots=[1,2],nodes=16, burst_intensity=10):
    matrix=np.zeros((nodes,nodes))
    for src in hotspots:
        for dst in range(nodes):
            if src != dst:
                matrix[src][dst] +=np.random.exponential(burst_intensity)

    # 剧烈热点流量
    hotspot=np.random.choice(hotspots)
    src=hotspot
    dst=np.random.choice(nodes)
    if src!=dst:
        matrix[src][dst]+=np.random.exponential(burst_intensity*5)

    return matrix

def generate_traffic_records(node=16,scale_factors=[1.0, 1.5 ,2.0,2.5, 3.0,3.5 ,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0]):
    # 包含所有流量记录的字典
    records = {
        "all_to_all":[],
        "burst":[],
        "ring_allreduce": [],
        "all_to_all_burst": [],
        "ring_allreduce_burst": []
    }
    # 生成all_to_all流量和burst流量和allreduce流量
    all_to_all=generate_mp_traffic_matrix(node,shard_nodes=[0,3,8,13],traffic_size=0.03125)
    records["all_to_all"].append({
        "matrix": all_to_all.tolist()
    })
    ring_allreduce=generate_ring_allreduce(node,traffic_size=44)
    records["ring_allreduce"].append({
        "matrix": ring_allreduce.tolist()
    })
    # burst_traffic = generate_burst_traffic(hotspots=[1, 2], nodes=16, burst_intensity=10)
    for scale in scale_factors:
        burst_intensity = scale * 10
        burst_traffic = generate_burst_traffic(hotspots=[1, 2], nodes=16, burst_intensity=burst_intensity)
        # 对生成的burst流量求和
        total_burst = 0
        for i in range(node):
            for j in range(node):
                if i != j:
                    total_burst += burst_traffic[i][j]
        all_to_all_burst=all_to_all + burst_traffic
        ring_allreduce_burst=ring_allreduce+burst_traffic
        records["burst"].append({
            "base_scale": scale,
            "burst_intensity":burst_intensity,
            "matrix": burst_traffic.tolist(),
            "total_burst": total_burst
        })
        records["all_to_all_burst"].append({
            "base_scale": scale,
            "burst_intensity": burst_intensity,
            "matrix": all_to_all_burst.tolist()
        })
        records["ring_allreduce_burst"].append({
            "base_scale": scale,
            "burst_intensity": burst_intensity,
            "matrix": ring_allreduce_burst.tolist()
        })
    return records

def plot_heatmap(matrix, save_path="mp_traffic_heatmap.png"):
    """
    绘制流量热点图
    :param matrix: 流量矩阵
    :param save_path: 图片保存路径
    """
    plt.figure(figsize=(12, 10))

    # 创建自定义颜色映射（突出显示零值）
    cmap = plt.cm.cividis
    cmap.set_under('black')  # 零值显示为黑色
    print(matrix.max())
    # 绘制热图
    img = plt.imshow(matrix,
                     cmap=cmap,
                     norm=colors.LogNorm(vmin=1, vmax=matrix.max()),  # 对数标准化
                     aspect='auto')

    # 添加坐标轴标签
    plt.title("AllReduce Traffic Matrix Heatmap", fontsize=14, pad=20)
    plt.xlabel("Receiver Node", fontsize=12)
    plt.ylabel("Sender Node", fontsize=12)

    # 设置坐标刻度
    max_nodes = matrix.shape[0]
    plt.xticks(np.arange(0, max_nodes, 1),
               labels=[f"S{i}" for i in range(0, max_nodes, 1)])
    plt.yticks(np.arange(0, max_nodes, 1),
               labels=[f"S{i}" for i in range(0, max_nodes, 1)])

    # 添加颜色条
    cbar = plt.colorbar(img, shrink=0.8)
    cbar.set_label('Traffic Size (GB)', rotation=270, labelpad=20)

    # 添加网格线
    plt.grid(which='minor', color='gray', linestyle=':', linewidth=0.5)

    # 保存并显示
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()


if __name__ == "__main__":
    # 生成示例矩阵（16节点，分片节点0,3,8,13）
    traffic_matrix = generate_mp_traffic_matrix(
        nodes=16,
        shard_nodes=[0, 3, 8, 13],
        traffic_size=32
    )
    traffic_matrix_1=generate_ring_allreduce(nodes=16,traffic_size=44)
    traffic_matrix_2=generate_burst_traffic(hotspots=[1,2],nodes=16,burst_intensity=10)
    # # 打印矩阵验证
    # print("Burst Traffic Matrix Sample:")
    # # print(traffic_matrix[0])  # 显示第一个分片节点的发送情况
    #
    # # 绘制热图
    # plot_heatmap(traffic_matrix_2)

    # 生成数据
    traffic_data = generate_traffic_records(node=16,scale_factors=[1.0, 1.5 ,2.0,2.5, 3.0,3.5 ,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0])
    # 保存为json文件
    with open("optimize_traffic_data.json", "w") as f:
        json.dump(traffic_data, f, indent=2)
    print("Traffic data saved to optimize_traffic_data.json")